# SimpleSequentialChain & SequentialChain
from langchain.chains.llm import LLMChain
from langchain.chains.sequential import SimpleSequentialChain
from langchain_community.chat_models import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
import __init__

chat_model = ChatOpenAI(
    temperature=0,
    model="gpt-3.5-turbo"
)

# chain 1
first_prompt = ChatPromptTemplate.from_template("帮我给{product}的公司起一个响亮容易记忆的名字？")
chain_one = LLMChain(
    llm=chat_model,
    prompt=first_prompt,
    verbose=True,
)

# chain 2
second_prompt = ChatPromptTemplate.from_template("用5个词来描述一下这个公司的名字：{company_name}")
chain_two = LLMChain(
    llm=chat_model,
    prompt=second_prompt,
    verbose=True,
)

overall_simple_chain = SimpleSequentialChain(
    chains=[chain_one, chain_two],
    verbose=True,  # 打开日志
)

overall_simple_chain.run("AI教育培训机构")
